
#include <string.h>
#include <stdlib.h>
#include "flash.h"

#define FLASH_W25Q08 	0XEF13
#define FLASH_W25Q16 	0XEF14
#define FLASH_W25Q32 	0XEF15
#define FLASH_W25Q64 	0XEF16
#define FLASH_W25Q128	0XEF17

#define FLASH_BASE_W25QXX           ((flash_type - FLASH_TYPE_W25QXX) * 0x01000000)
#define FLASH_REAL_ADDR_W25QXX(x)   (x - (flash_type - FLASH_TYPE_W25QXX) * 0x01000000)
#define FLASH_SECTOR_SIZE_W25QXX    4096

#define W25X_WriteEnable		0x06
#define W25X_WriteDisable		0x04
#define W25X_ReadStatusReg		0x05
#define W25X_WriteStatusReg		0x01
#define W25X_ReadData			0x03
#define W25X_FastReadData		0x0B
#define W25X_FastReadDual		0x3B
#define W25X_PageProgram		0x02
#define W25X_BlockErase			0xD8
#define W25X_SectorErase		0x20
#define W25X_ChipErase			0xC7
#define W25X_PowerDown			0xB9
#define W25X_ReleasePowerDown	0xAB
#define W25X_DeviceID			0xAB
#define W25X_ManufactDeviceID	0x90
#define W25X_JedecDeviceID		0x9F

static uint8_t flash_read_sr_w25qxx(flash_type_e flash_type) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    uint8_t data = W25X_ReadStatusReg;

    spi->take(flash->cfg.cs_pin);
    spi->write_read(&data, NULL, 1);
    data = 0xFF;
    spi->write_read(&data, &data, 1);
    spi->release(flash->cfg.cs_pin);

    return data;
}

#if 0
static void flash_write_sr_w25qxx(flash_type_e flash_type, uint8_t sr) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    uint8_t data = W25X_WriteStatusReg;

    spi->take(flash->cfg.cs_pin);
    spi->write_read(&data, NULL, 1);
    spi->write_read(&sr, NULL, 1);
    spi->release(flash->cfg.cs_pin);
}
#endif

static void flash_write_enable_w25qxx(flash_type_e flash_type) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    uint8_t data = W25X_WriteEnable;

    spi->take(flash->cfg.cs_pin);
    spi->write_read(&data, NULL, 1);
    spi->release(flash->cfg.cs_pin);
}

static void flash_write_disable_w25qxx(flash_type_e flash_type) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    uint8_t data = W25X_WriteDisable;

    spi->take(flash->cfg.cs_pin);
    spi->write_read(&data, NULL, 1);
    spi->release(flash->cfg.cs_pin);
}

static uint16_t flash_read_id_w25qxx(flash_type_e flash_type) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    uint8_t data[6] = {W25X_ManufactDeviceID, 0x00, 0x00, 0x00, 0xFF, 0xFF};

    spi->take(flash->cfg.cs_pin);
    spi->write_read(data, data, sizeof(data));
    spi->release(flash->cfg.cs_pin);

    return ((uint16_t)data[4] << 8) | data[5];
}

static uint32_t flash_get_size_w25qxx(flash_type_e flash_type) {
    switch(flash_read_id_w25qxx(flash_type)) {
        case FLASH_W25Q08: return 1 * 1024 * 1024;
        case FLASH_W25Q16: return 2 * 1024 * 1024;
        case FLASH_W25Q32: return 4 * 1024 * 1024;
        case FLASH_W25Q64: return 8 * 1024 * 1024;
        case FLASH_W25Q128: return 16 * 1024 * 1024;
    }

    return (uint32_t)-1;
}

static void flash_busy_wait_w25qxx(flash_type_e flash_type) {
    while((flash_read_sr_w25qxx(flash_type) & 0x01) == 0x01);
}

static void flash_write_page_w25qxx(flash_type_e flash_type, uint32_t address, uint8_t *data, uint32_t len) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    address = FLASH_REAL_ADDR_W25QXX(address);
    uint8_t cmd[4] = {W25X_PageProgram, (uint8_t)(address >> 16), (uint8_t)(address >> 8), (uint8_t)address};

    flash_busy_wait_w25qxx(flash_type);
    spi->take(flash->cfg.cs_pin);
    spi->write_read(cmd, NULL, sizeof(cmd));
    spi->write_read(data, NULL, len);
    spi->release(flash->cfg.cs_pin);
    flash_busy_wait_w25qxx(flash_type);
}

static void flash_write_no_check_w25qxx(flash_type_e flash_type, uint32_t address, uint8_t *data, uint32_t len) {
    uint16_t page_remain;
	page_remain = 256 - address % 256;
	if(len <= page_remain) page_remain = len;
	while(1)
	{	   
		flash_write_page_w25qxx(flash_type, address, data, page_remain);
		if(len == page_remain) break;
	 	else
		{
			data += page_remain;
			address += page_remain;	

			len -= page_remain;
			if(len > 256) page_remain = 256;
			else page_remain = len;
		}
	};
}

static void flash_erase_sector_w25qxx(flash_type_e flash_type, uint32_t sec_pos) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    uint32_t sec_addr = sec_pos * FLASH_SECTOR_SIZE_W25QXX;
    uint8_t cmd[4] = {W25X_SectorErase, (uint8_t)(sec_addr >> 16), (uint8_t)(sec_addr >> 8), (uint8_t)sec_addr};

    flash_busy_wait_w25qxx(flash_type);
    spi->take(flash->cfg.cs_pin);
    spi->write_read(cmd, NULL, sizeof(cmd));
    spi->release(flash->cfg.cs_pin);
    flash_busy_wait_w25qxx(flash_type);
}

int flash_read_w25qxx(flash_type_e flash_type, uint32_t address, uint8_t *data, uint32_t len) {
    uint32_t flash_size = flash_get_size_w25qxx(flash_type);
    if(flash_size == (uint32_t)-1) return -1;
	if(address < FLASH_BASE_W25QXX || (address >= (FLASH_BASE_W25QXX + flash_size))) return -1;

    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    address = FLASH_REAL_ADDR_W25QXX(address);
    uint8_t cmd[4] = {W25X_ReadData, (uint8_t)(address >> 16), (uint8_t)(address >> 8), (uint8_t)address};
    uint8_t dummy = 0xFF;

    spi->take(flash->cfg.cs_pin);
    spi->write_read(cmd, NULL, sizeof(cmd));
    for(int i=0; i<len; i++) {
        spi->write_read(&dummy, data+i, 1);
    }
    spi->release(flash->cfg.cs_pin);

    return 0;
}

int flash_write_w25qxx(flash_type_e flash_type, uint32_t address, uint8_t *data, uint32_t len) {
	uint32_t start_addr = address - FLASH_BASE_W25QXX;
	uint32_t end_addr = address - FLASH_BASE_W25QXX + len;
	uint32_t start_sector_pos = start_addr / FLASH_SECTOR_SIZE_W25QXX;
	uint32_t end_sector_pos = end_addr / FLASH_SECTOR_SIZE_W25QXX;
	uint32_t start_wbyte_pos_insector = 0;
	uint32_t end_wbyte_pos_insector = 0;
	uint8_t *flash_rw_buf;
    uint32_t flash_size = flash_get_size_w25qxx(flash_type);

    if(flash_size == (uint32_t)-1) return -1;
	if(address < FLASH_BASE_W25QXX || (address >= (FLASH_BASE_W25QXX + flash_size))) return -1;

	flash_write_enable_w25qxx(flash_type);

	flash_rw_buf = malloc(FLASH_SECTOR_SIZE_W25QXX);
	if(!flash_rw_buf) {
		flash_write_disable_w25qxx(flash_type);
		return -1;
	}

	for(uint32_t pos = start_sector_pos; pos <= end_sector_pos; pos++) {
		int i;

		if(pos == start_sector_pos) start_wbyte_pos_insector = start_addr % FLASH_SECTOR_SIZE_W25QXX;
		else start_wbyte_pos_insector = 0;

		if(pos == end_sector_pos) end_wbyte_pos_insector = end_addr % FLASH_SECTOR_SIZE_W25QXX;
		else end_wbyte_pos_insector = FLASH_SECTOR_SIZE_W25QXX;

		flash_read_w25qxx(flash_type, pos * FLASH_SECTOR_SIZE_W25QXX + FLASH_BASE_W25QXX, flash_rw_buf, FLASH_SECTOR_SIZE_W25QXX);

		if(memcmp(flash_rw_buf + start_wbyte_pos_insector, data, end_wbyte_pos_insector - start_wbyte_pos_insector) == 0)
			goto next_loop;

		for(i=0; i<FLASH_SECTOR_SIZE_W25QXX; i++) {
			if(flash_rw_buf[i] != 0xFF) break;
		}

		if(i < FLASH_SECTOR_SIZE_W25QXX) {
            flash_erase_sector_w25qxx(flash_type, pos);
		}

		memcpy(flash_rw_buf + start_wbyte_pos_insector, data, end_wbyte_pos_insector - start_wbyte_pos_insector);
		flash_write_no_check_w25qxx(flash_type, pos * FLASH_SECTOR_SIZE_W25QXX + FLASH_BASE_W25QXX, flash_rw_buf, FLASH_SECTOR_SIZE_W25QXX);

	next_loop:
		data += end_wbyte_pos_insector - start_wbyte_pos_insector;
	}

	flash_write_disable_w25qxx(flash_type);
	free(flash_rw_buf);
	return 0;
}

int flash_erase_w25qxx(flash_type_e flash_type, uint32_t address, uint32_t len) {
	uint32_t start_addr = address - FLASH_BASE_W25QXX;
	uint32_t end_addr = address - FLASH_BASE_W25QXX + len;
	uint32_t start_sector_pos = start_addr / FLASH_SECTOR_SIZE_W25QXX;
	uint32_t end_sector_pos = end_addr / FLASH_SECTOR_SIZE_W25QXX;
	uint8_t *flash_rw_buf;
    uint32_t flash_size = flash_get_size_w25qxx(flash_type);

    if(flash_size == (uint32_t)-1) return -1;
	if(address < FLASH_BASE_W25QXX || (address >= (FLASH_BASE_W25QXX + flash_size))) return -1;

	flash_write_enable_w25qxx(flash_type);

	flash_rw_buf = malloc(FLASH_SECTOR_SIZE_W25QXX);
	if(!flash_rw_buf) {
		flash_write_disable_w25qxx(flash_type);
		return -1;
	}

	for(uint32_t pos = start_sector_pos; pos <= end_sector_pos; pos++) {
		int i;

		flash_read_w25qxx(flash_type, pos * FLASH_SECTOR_SIZE_W25QXX + FLASH_BASE_W25QXX, flash_rw_buf, FLASH_SECTOR_SIZE_W25QXX);
		for(i=0; i<FLASH_SECTOR_SIZE_W25QXX; i++) {
			if(flash_rw_buf[i] != 0xFF) break;
		}

		if(i < FLASH_SECTOR_SIZE_W25QXX) {
            flash_erase_sector_w25qxx(flash_type, pos);
		}
	}

	flash_write_disable_w25qxx(flash_type);
	free(flash_rw_buf);
	return 0;
}

int flash_open_w25qxx(flash_type_e flash_type, flash_cfg_t *config) {
    if(!config) return -1;

    if(config->spi_port >= SPI_PORT_MAX) return -1;
    if(config->cs_pin >= GPIO_PIN_MAX) return -1;

    spi_t *spi = spi_find_by_port(config->spi_port);
    if(!spi) return -1;

    spi_cfg_t spi_cfg = {
        .dir = SPI_DIR_2LINES_FULL_DUPLEX,
        .mode = SPI_MODE_MASTER,
        .cpol = SPI_CPOL_HIGH,
        .cpha = SPI_CPHA_2EDGE,
        .freq = SPI_FREQ_18M,
        .first_bit = SPI_1ST_BIT_MSB
    };

    if(spi->add_cs_pin(config->cs_pin)) return -1;
    if(spi->open(&spi_cfg)) return -1;

    flash_t *flash = flash_find_by_type(flash_type);
    memcpy(&flash->cfg, config, sizeof(flash_cfg_t));
    flash->info.start_addr = FLASH_BASE_W25QXX;
    flash->info.size = flash_get_size_w25qxx(flash_type);
    flash->info.block_size = FLASH_SECTOR_SIZE_W25QXX;
    return 0;
}

void flash_close_w25qxx(flash_type_e flash_type) {
    flash_t *flash = flash_find_by_type(flash_type);
    spi_t *spi = spi_find_by_port(flash->cfg.spi_port);
    spi->close();
}
